
#include <string.h>
#include <stddef.h>
#include "ws_client.h"

// sec-websocket-key  src:gilfoylehandsome
static const char *_ws_key = "Z2lsZm95bGVoYW5kc29tZQ==";
// mask key
static const char _mask_key[4] = {0x37, 0xfa, 0x21, 0x3d};

size_t ws_build_handshake_req(char *buf, 
		size_t buf_sz,
		char const *ip,
		char const *port,
		char const *path)
{
	if (buf == NULL || buf_sz < 128) return 0;

	size_t offset = 0;

	// first line
	strcpy(buf+offset, "GET ");
	offset = strlen(buf);
	strcpy(buf+offset, path);
	offset = strlen(buf);
	strcpy(buf+offset, " HTTP/1.1\r\n");
	offset = strlen(buf);

	// HOST
	strcpy(buf+offset, "HOST: ");
	offset = strlen(buf);
	strcpy(buf+offset, ip);
	offset = strlen(buf);
	buf[offset++] = ':';
	strcpy(buf+offset, port);
	offset = strlen(buf);
	strcpy(buf+offset, "\r\n");
	offset = strlen(buf);

	// Connection
	strcpy(buf+offset, "Connection: Upgrade\r\n");
	offset = strlen(buf);

	// Upgrade
	strcpy(buf+offset, "Upgrade: websocket\r\n");
	offset = strlen(buf);

	// Origin
	strcpy(buf+offset, "Origin: null\r\n");
	offset = strlen(buf);

	// Sec-WebSocket-Version
	strcpy(buf+offset, "Sec-WebSocket-Version: 13\r\n");
	offset = strlen(buf);

	// Sec-WebSocket-Key
	strcpy(buf+offset, "Sec-WebSocket-Key: ");
	offset = strlen(buf);
	strcpy(buf+offset, _ws_key);
	offset = strlen(buf);

	strcpy(buf+offset, "\r\n\r\n");
	offset = strlen(buf);

	return offset;
}

bool ws_check_handshake_rsp(char const *buf)
{
	return (strstr(buf, "101")) == NULL ? false : true; 
}

size_t ws_build_frame(struct ws_frame *frame, char *buf, size_t buf_sz)
{
	if (frame == NULL || buf == NULL) return 0;

	int head_len = 0, offset = 0;

	head_len = (frame->payloadlen<0x7e) ? 6 : 8;

	if (buf_sz < frame->payloadlen + head_len) 
	{
		/* buf_sz too small */
		return 0;
	}

	buf[offset++] = 0x80 | frame->opcode;

	if (head_len == 6) 
	{
		buf[offset++] = 0x80 | frame->payloadlen;
	}
	else if (head_len == 8) 
	{
		buf[offset++] = 0x80 | 0x7e;
		buf[offset++] = (frame->payloadlen>>8) & 0xff;
		buf[offset++] = frame->payloadlen & 0xff;
	}

	buf[offset++] = _mask_key[0]; 
	buf[offset++] = _mask_key[1];
	buf[offset++] = _mask_key[2];
	buf[offset++] = _mask_key[3];

	if (frame->payloadlen > 0)
	{
		memcpy(buf+offset, frame->payload, frame->payloadlen);
	}

	/* mask payload data */
	for(int i=0; i<frame->payloadlen; i++) 
	{
		buf[offset] = buf[offset] ^ (_mask_key[i%4]);
		offset++;
	}

	return offset;
}

bool ws_parse_recv_frame(struct ws_frame *frame, const char *buf, size_t buf_sz)
{
	if (frame == NULL || buf == NULL || buf_sz < 2) return false;

	size_t payload_len = 0, head_len = 0;
	char mask_flag = 0;
	char mask[4] = {0};

	/* first byte */
	if ((buf[0] & 0x80) != 0x80) 
	{
		/* dont parse split pkg */
		return false;
	}
	frame->opcode = buf[0] & 0x0f;
	head_len++;

	/*second byte */
	mask_flag = (buf[1]&0x80) == 0x80;
	payload_len = buf[1] & 0x7f;
	head_len++;

	/* get payload len */
	if (payload_len < 0x7e) 
	{
		payload_len = payload_len;
	}
	else if (payload_len == 0x7e) 
	{	  		
		payload_len = (buf[2] & 0xff) << 8 | (buf[3] & 0xff);
		head_len += 2;
	}
	else 
	{
		return false;
	}

	/* if has mask */
	if (mask_flag) 
	{
		head_len += 4;
	}

	/* verify total len */
	if (buf_sz < (head_len + payload_len)) 
	{
		return false;
	}

	/* get payloadlen and payload */
	frame->payloadlen = payload_len;
	if (payload_len > 0) 
	{
		memcpy(frame->payload, buf+head_len, payload_len);
	}
	else 
	{
		frame->payload = NULL;
	}

	/* mask payload data */
	if (mask_flag) 
	{
		if (payload_len < 0x7e) 
		{
			memcpy(mask, buf+2, 4);
		}
		else 
		{
			memcpy(mask, buf+4, 4);
		}

		for (int i=0; i<payload_len; i++) 
		{
			frame->payload[i] = frame->payload[i] ^ mask[i%4];
		}
	}

	return true;
}

